import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import MinMaxScaler


def create_dataset(dataset, lookback):
    X, y = [], []
    for i in range(len(dataset) - lookback):
        feature = dataset[i:i + lookback, :-1]
        target = dataset[i + lookback - 1, -1]
        X.append(feature)
        y.append(target)
    return np.array(X), np.array(y)


def create_train_data(lookback, target_column='minTemp'):
    df = pd.read_csv("predicte/formatted_data.csv")
    df = df[['DATE', 'maxTemp', 'minTemp', 'weather_type']]
    df['DATE'] = pd.to_datetime(df['DATE'])
    df['year'] = df['DATE'].dt.year
    df['Day_of_Week'] = df['DATE'].dt.dayofweek
    df['month'] = df['DATE'].dt.month
    df['day'] = df['DATE'].dt.day
    df = pd.get_dummies(df, columns=['weather_type'], drop_first=True)

    # Shift features except the target column
    cols = [col for col in df.columns if col not in ['DATE', 'minTemp']]
    df_shifted = df[cols].shift(1)
    df = pd.concat([df[['DATE', 'minTemp']], df_shifted], axis=1)
    df = df.dropna()

    # Define features without 'maxTemp'
    features = ['maxTemp'] + [col for col in df.columns if 'weather_type_' in col] + ['month', 'day', 'Day_of_Week']
    print(features)

    scaler = MinMaxScaler()
    df_scaled = scaler.fit_transform(df[features].astype('float32'))

    # Add the target column (maxTemp) back for creating datasets
    df_scaled_with_target = np.hstack((df_scaled, df['minTemp'].values.reshape(-1, 1).astype('float32')))

    X_lookback, y_lookback = create_dataset(df_scaled_with_target, lookback=lookback)
    y_lookback = y_lookback.reshape(-1, 1)

    df_original_dates = pd.DataFrame(df['DATE'].values, columns=['DATE'])
    df_original_dates = df_original_dates.iloc[lookback - 1:]
    idx_test = df_original_dates.loc[df_original_dates['DATE'] == '2024-02-04'].index[0]
    idx_val = df_original_dates.loc[df_original_dates['DATE'] == '2024-10-30'].index[0]

    df.to_csv('new_data1.csv', index=False)

    X_train, y_train = torch.tensor(X_lookback[:(idx_val - lookback)]), torch.tensor(y_lookback[:(idx_val - lookback)])
    X_val, y_val = torch.tensor(X_lookback[(idx_val - lookback):(idx_test - lookback)]), torch.tensor(
        y_lookback[(idx_val - lookback):(idx_test - lookback)])
    X_test, y_test = torch.tensor(X_lookback[(idx_test - lookback):]), torch.tensor(y_lookback[(idx_test - lookback):])

    return X_train, y_train, X_val, y_val, X_test, y_test, scaler
